#include "CLibssh2.h"

using namespace std;
#ifdef WIN32
#include <winsock2.h>
#include <windows.h>
#include <sys/stat.h>
#else
#include <arpa/inet.h>
#include <fcntl.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#endif // WIN32

#include "libssh2.h"


static  char sg_password[16];
static void kbd_callback(const char *name, int name_len,
                         const char *instruction, int instruction_len,
                         int num_prompts,
                         const LIBSSH2_USERAUTH_KBDINT_PROMPT *prompts,
                         LIBSSH2_USERAUTH_KBDINT_RESPONSE *responses,
                         void **abstract)
{
    (void)name;
    (void)name_len;
    (void)instruction;
    (void)instruction_len;
    if (num_prompts == 1)
    {
        responses[0].text = strdup(sg_password);
        responses[0].length = strlen(sg_password);
    }
    (void)prompts;
    (void)abstract;
} /* kbd_callback */

void CLibssh2::init()
{

    int errcode = libssh2_init(0);
    if (errcode != 0)
        return;
}

void CLibssh2::fini()
{
    libssh2_exit();
}

CLibssh2::CLibssh2(const std::string& ip, uint16_t port, const std::string& username, const std::string& password, uint32_t timeout_seconds, bool nonblocking)
    : _socket_fd(-1), _session(NULL), _ip(ip), _port(port), _username(username), _timeout_seconds(timeout_seconds)
{
    int re;
    re = create_session(nonblocking);
    if(re == -1)
    {
        link_mark = -1;
        return ;
    }
    re = handshake();
    if(re == -1)
    {
        link_mark = -1;
        return ;
    }
    re = set_known_hosts();
    if(re == -1)
    {
        link_mark = -1;
        return ;
    }
    re = validate_authorization(password);
    if(re == -1)
    {
        link_mark = -1;
        return ;
    }
}

CLibssh2::~CLibssh2()
{
#ifdef WIN32
    WSACleanup();
#endif
    cleanup();
}

void CLibssh2::my_u_sleep(int time)
{
#ifdef WIN32
    Sleep(time/10000);
#else
    usleep(time);
#endif
}

void CLibssh2::remotely_execute(
    const std::string& command, std::ostream& out,
    int* exitcode, std::string* exitsignal, std::string* errmsg, int64_t* num_bytes)
{
    LIBSSH2_CHANNEL* channel = static_cast<LIBSSH2_CHANNEL*>(open_ssh_channel());

    while (true)
    {
        int errcode = libssh2_channel_exec(channel, command.c_str());
        if (0 == errcode)
        {
            break;
        }
        else if (errcode != LIBSSH2_ERROR_EAGAIN)
        {
            *exitcode = close_ssh_channel(channel, exitsignal, errmsg);
            return;
        }
        else
        {
            my_u_sleep(10);
        }
    }
    *num_bytes = read_channel(channel, out, NULL);
    *exitcode = close_ssh_channel(channel, exitsignal, errmsg);
}

void CLibssh2::download(const std::string& remote_filepath, std::ostream& out, int64_t* num_bytes)
{
    struct stat fileinfo;
    LIBSSH2_CHANNEL* channel = static_cast<LIBSSH2_CHANNEL*>(open_scp_read_channel(remote_filepath, &fileinfo));
    *num_bytes = read_channel(channel, out, &fileinfo);
    libssh2_channel_free(channel);
}

void CLibssh2::upload(const std::string& local_filepath, const std::string& remote_filepath, int64_t* num_bytes)
{
    int local_fd = -1;
    LIBSSH2_CHANNEL* channel = NULL;


    struct stat fileinfo;

    local_fd = open(local_filepath.c_str(), 0);
    if (-1 == local_fd)
    {
        return;
    }
    if (-1 == fstat(local_fd, &fileinfo))
    {
        return;
    }
    if (0 == fileinfo.st_size)
    {
        return;
    }
    close(local_fd);

    channel = static_cast<LIBSSH2_CHANNEL*>(open_scp_write_channel(
            remote_filepath, fileinfo.st_mode, fileinfo.st_size, fileinfo.st_mtime, fileinfo.st_atime));

    *num_bytes = 0;
    ifstream local_in(local_filepath.c_str(), ios::in|ios::binary|ios::ate);
    if(!local_in)
    {
        libssh2_channel_send_eof(channel); // ����Զ�������ѷ�����ϣ������⼸�����Զ��յ����ֽ������ܲ��㹻��
        libssh2_channel_wait_eof(channel); // �ȴ�Զ�˻�Ӧlibssh2_channel_send_eof()
        libssh2_channel_wait_closed(channel); // �ȴ�Զ�˹ر�ͨ��
        libssh2_channel_free(channel);
        return;
    }
    *num_bytes = 0;
    local_in.seekg(0, std::ios::beg);
    while(!local_in.eof())
    {
        char buffer[4096] = {0};
        int bytes = 0;
        local_in.read(buffer,sizeof(buffer));
        bytes = local_in.gcount();
        *num_bytes += bytes;
        write_channel(channel, buffer, bytes);
    }
    local_in.close();
    libssh2_channel_send_eof(channel); // ����Զ�������ѷ�����ϣ������⼸�����Զ��յ����ֽ������ܲ��㹻��
    libssh2_channel_wait_eof(channel); // �ȴ�Զ�˻�Ӧlibssh2_channel_send_eof()
    libssh2_channel_wait_closed(channel); // �ȴ�Զ�˹ر�ͨ��
    libssh2_channel_free(channel);

}

int CLibssh2::get_session_errcode() const
{
    LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);
    return libssh2_session_last_errno(session);
}

std::string CLibssh2::get_session_errmsg() const
{
    std::string result;
    char* errmsg;
    int errmsg_len = 0;
    LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);

    (void)libssh2_session_last_error(session, &errmsg, &errmsg_len, 1);
    if (errmsg != NULL)
    {
        result = errmsg;
        free(errmsg);
    }

    return result;
}

void CLibssh2::cleanup()
{
    if (_session != NULL)
    {
        LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);
        libssh2_session_disconnect(session, "normal Shutdown");
        libssh2_session_free(session);
        _session = NULL;
    }
    if (_socket_fd != -1)
    {
#ifdef WIN32
        closesocket(_socket_fd);
#else
        close(_socket_fd);
#endif
        _socket_fd = -1;
    }
}

int  CLibssh2::create_session(bool nonblocking)
{
#ifdef WIN32
    WSADATA wsaData;
    int nRet;
    if((nRet = WSAStartup(MAKEWORD(2,2),&wsaData)) != 0)
    {
        return -1;
    }
#endif // WIN32
    int socket_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (-1 == socket_fd)
    {
#ifdef WIN32
        closesocket(_socket_fd);
#else
        close(_socket_fd);
#endif
        return -1;
    }

    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(_port);
    addr.sin_addr.s_addr = inet_addr(_ip.c_str());

    if (-1 == connect(socket_fd, (struct sockaddr*)&addr, sizeof(struct sockaddr_in)))
    {
#ifdef WIN32
        closesocket(_socket_fd);
#else
        close(_socket_fd);
#endif
        //throw -1;
        return -1;
    }

    // ����ssh2�Ự
    LIBSSH2_SESSION* session = libssh2_session_init();
    if (nonblocking) // ����Ϊ������
        libssh2_session_set_blocking(session, 0);
    else if (_timeout_seconds > 0)
        libssh2_session_set_timeout(session, _timeout_seconds*1000);

    _session = session;
    _socket_fd = socket_fd;
    return 0;

}

int CLibssh2::set_known_hosts()
{
    int type;
    int check;
    size_t len;
    struct libssh2_knownhost* known_host;
    LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);
    LIBSSH2_KNOWNHOSTS* known_hosts = libssh2_knownhost_init(session);

    if (NULL == known_hosts)
    {
        libssh2_knownhost_free(known_hosts);
        return -1;
    }


    // read all hosts from here
    libssh2_knownhost_readfile(known_hosts, "known_hosts", LIBSSH2_KNOWNHOST_FILE_OPENSSH);
    // store all known hosts to here
    libssh2_knownhost_writefile(known_hosts, "dumpfile", LIBSSH2_KNOWNHOST_FILE_OPENSSH);

    const char* fingerprint = libssh2_session_hostkey(session, &len, &type);
    if (NULL == fingerprint)
    {
        libssh2_knownhost_free(known_hosts);
        return -1;
    }

#if LIBSSH2_VERSION_NUM >= 0x010206
    // introduced in 1.2.6
    check = libssh2_knownhost_checkp(
                known_hosts, _ip.c_str(), _port,
                fingerprint, len, LIBSSH2_KNOWNHOST_TYPE_PLAIN|LIBSSH2_KNOWNHOST_KEYENC_RAW, &known_host);
#else
    // 1.2.5 or older
    check = libssh2_knownhost_check(
                known_hosts, _ip.c_str(),
                fingerprint, len, LIBSSH2_KNOWNHOST_TYPE_PLAIN|LIBSSH2_KNOWNHOST_KEYENC_RAW, &known_host);
#endif // LIBSSH2_VERSION_NUM

    if (check <= LIBSSH2_KNOWNHOST_CHECK_MISMATCH)
    {
        // At this point, we could verify that 'check' tells us the key is fine or bail out.
    }
    libssh2_knownhost_free(known_hosts);
    return 0;

}

int CLibssh2::validate_authorization(const std::string& password)
{
    int errcode;
    LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);

    while (true)
    {
        // ��ѯ�����֧�ֵ���Ȩ��ʽ�б�
        const char* userauth_list = libssh2_userauth_list(session, _username.c_str(), _username.size());
        int userauth = 0;

        if (NULL == userauth_list)
        {
            errcode = get_session_errcode();
            if (LIBSSH2_ERROR_EAGAIN == errcode)
            {
                continue;
            }
            else
            {
                return -1;
            }
        }

        if (strstr(userauth_list, "password") != NULL)
        {
            // ���뷽ʽ��
            // Ҫ�����˵������ļ�/etc/ssh/sshd_config��PasswordAuthenticationֵΪyes
            userauth |= 0x01;
        }
        if (strstr(userauth_list, "keyboard-interactive") != NULL)
        {
            // ������ʽ
            userauth |= 0x02;
            strncpy(sg_password, password.c_str(), sizeof(sg_password)-2);
            sg_password[sizeof(sg_password)-1] = '\0';
        }
        if (strstr(userauth_list, "publickey") != NULL)
        {
            userauth |= 0x04;
        }

        errcode = -1;
        if (userauth & 0x01) // password
        {
            errcode = libssh2_userauth_password(session, _username.c_str(), password.c_str());
        }
        if (0 == errcode)
        {
            break;
        }
        if ((userauth & 0x02) && (errcode != LIBSSH2_ERROR_EAGAIN)) // keyboard
        {
            errcode = libssh2_userauth_keyboard_interactive(session, _username.c_str(), &kbd_callback);
        }
        if (0 == errcode)
        {
            break;
        }
        if ((userauth & 0x04) && (errcode != LIBSSH2_ERROR_EAGAIN)) // publickey
        {
            errcode = libssh2_userauth_publickey_fromfile(session, _username.c_str(),
                      "~/.ssh/id_rsa.pub", "~/.ssh/id_rsa", password.c_str());
        }
        if (0 == errcode)
        {
            break;
        }

        if (-1 == errcode)
        {
            return -1;
        }
        else if (errcode != LIBSSH2_ERROR_EAGAIN)
        {
            return -1;
        }
    }
    return 0;
}


int CLibssh2::handshake()
{
    LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);

    while (true)
    {
        // ssh2����
        int errcode = libssh2_session_handshake(session, _socket_fd);
        if (0 == errcode)
        {
            break;
        }
        else if (errcode != LIBSSH2_ERROR_EAGAIN)
        {
            return -1;
        }
    }
    return 0;
}

void* CLibssh2::open_ssh_channel()
{
    LIBSSH2_CHANNEL* channel = NULL;
    LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);

    // �����Ựͨ��
    while (true)
    {
        channel = libssh2_channel_open_session(session);
        if (channel != NULL)
            break;

        int errcode = get_session_errcode();
        if (errcode != LIBSSH2_ERROR_EAGAIN)
        {
            return NULL;
        }
    }

    return channel;
}

void* CLibssh2::open_scp_read_channel(const std::string& remote_filepath, struct stat* fileinfo)
{
    LIBSSH2_CHANNEL* channel = NULL;
    LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);

    for (;;)
    {
        channel = libssh2_scp_recv(session, remote_filepath.c_str(), fileinfo);
        if (channel != NULL)
            break;

        int errcode = get_session_errcode();
        if (errcode != LIBSSH2_ERROR_EAGAIN)
        {
            return NULL;
        }
    }

    return channel;
}

void* CLibssh2::open_scp_write_channel(const std::string& remote_filepath, int filemode, size_t filesize, time_t mtime, time_t atime)
{
    LIBSSH2_CHANNEL* channel = NULL;
    LIBSSH2_SESSION* session = static_cast<LIBSSH2_SESSION*>(_session);

    for (;;)
    {
        //channel = libssh2_scp_send_ex(session, remote_filepath.c_str(), filemode&0777, filesize, (long)mtime, (long)atime);
        channel = libssh2_scp_send(session, remote_filepath.c_str(), filemode&0777, filesize);
        if (channel != NULL)
            break;

        int errcode = get_session_errcode();
        if (errcode != LIBSSH2_ERROR_EAGAIN)
        {
            return NULL;
        }
    }

    return channel;
}

int CLibssh2::close_ssh_channel(void* channel, std::string* exitsignal, std::string* errmsg)
{
    int exitcode = 127; // 127: command not exists
    LIBSSH2_CHANNEL* channel_ = static_cast<LIBSSH2_CHANNEL*>(channel);

    while (channel_ != NULL)
    {
        int errcode = libssh2_channel_close(channel_);
        if (0 == errcode)
        {
            char* errmsg_ = NULL;
            char* exitsignal_ = NULL;

            exitcode = libssh2_channel_get_exit_status(channel_);
            if (exitcode != 0) // 0 success
            {
                // ���ö˿���strerror(*exitcode)ȡ�ó���ԭ��
                libssh2_channel_get_exit_signal(channel_, &exitsignal_, NULL, &errmsg_, NULL, NULL, NULL);
                if (errmsg_ != NULL)
                {
                    *errmsg = errmsg_;
                    free(errmsg_);
                }
                if (exitsignal_ != NULL)
                {
                    *exitsignal = exitsignal_;
                    free(exitsignal_);
                }
            }

            libssh2_channel_free(channel_);
            channel_ = NULL;
            break;
        }
        else if (LIBSSH2_ERROR_EAGAIN == errcode)
        {
            return exitcode;
        }
        else
        {
            return exitcode;
        }
    }

    return exitcode;
}

int64_t CLibssh2::read_channel(void* channel, std::ostream& out, const struct stat* fileinfo)
{
    int64_t num_bytes = 0;
    LIBSSH2_CHANNEL* channel_ = static_cast<LIBSSH2_CHANNEL*>(channel);

    while (true)
    {
        char buffer[4096];
        int amount = sizeof(buffer) - 1;
        if (fileinfo != NULL)
        {
            if ((fileinfo->st_size - num_bytes) < amount)
                amount = static_cast<int>(fileinfo->st_size - num_bytes);
        }

        const int bytes = libssh2_channel_read(channel_, buffer, amount);
        if (0 == bytes)
        {
            break; // connection closed now
        }
        else if (bytes > 0)
        {
            num_bytes += bytes;
            buffer[bytes] = '\0';
            out.write(buffer, bytes); // out << std::string(buffer, bytes);

            //С��Ҳ�������ģ��Ƿ�����ɡ�if (0 == bytes)������
            //if (bytes < static_cast<int>(sizeof(buffer)-1))
            //    break;
            if ((fileinfo != NULL) && (num_bytes >= fileinfo->st_size))
            {
                out.flush();
                break;
            }
        }
        else
        {
            int errcode = get_session_errcode();

            if (errcode != LIBSSH2_ERROR_EAGAIN)
            {
                return 0;
            }
        }
    }

    return num_bytes;
}

void CLibssh2::write_channel(void* channel, const char *buffer, size_t buffer_size)
{
    const char* buffer_ = buffer;
    size_t buffer_size_ = buffer_size;
    LIBSSH2_CHANNEL* channel_ = static_cast<LIBSSH2_CHANNEL*>(channel);

    for (;;)
    {
        int bytes = libssh2_channel_write(channel_, buffer_, buffer_size_);

        if (bytes == static_cast<int>(buffer_size_))
        {
            break;
        }
        else if (bytes > 0)
        {
            buffer_ += bytes;
            buffer_size_ -= bytes;
        }
        else
        {
            int errcode = get_session_errcode();

            if (errcode != LIBSSH2_ERROR_EAGAIN)
            {
                return;
            }
        }
    }
}
